# -*- coding: UTF-8 -*-

from neo4j import GraphDatabase
import math
import mashups
import web_apis
import user

RANDWALK_QUERY = '''
        CALL gds.alpha.randomWalk.stream({{
            nodeProjection: '*',
            relationshipProjection: {{
                    LINKS: {{
                        type: '*',
                        orientation: '{reverse}'
                    }}
                }},
            start: {start},
            steps: {steps},
            walks: {walks}
        }})
        YIELD nodeIds
        RETURN nodeIds
        '''

GET_NEXT_QUERY = '''
        MATCH (n)-[]->(t)
        WHERE id(n)={id}
        RETURN id(t)
        '''

GET_INDEGREE_QUERY = '''
        MATCH (t)-[]->(n)
        WHERE id(n)={id}
        RETURN count(t)
        '''

CREATE_CONSTRAINT_QUERY = '''
        CREATE CONSTRAINT ON (n:{data_type}) ASSERT n.{unique_attribute} IS UNIQUE
        '''

CREATE_NODE_QUERY = '''
        CREATE (n:{data_type}{{name:'{data_name}'}}) return id(n)
        '''

CREATE_RELATIONSHIP_QUERY = '''
        MATCH (from:{from_type}{{name:'{from_name}'}})
        MATCH (to:{to_type}{{name:'{to_name}'}})
        MERGE (from)-[:{relation_type}]->(to)
        '''

GET_NODE_QUERY = '''
        MATCH (n:{data_type}{{name:'{data_name}'}})
        return id(n)
        '''

# 删除API的边，再去预测
REMOVE_API_INVOKE_QUERY = '''    
        MATCH (:Mashup{{name:'{mashup_name}'}})-[e:INVOKE]->(a:API)
        DELETE (e)
        RETURN id(a)
        '''

# 测试完，再恢复这些边
RECOVER_API_INVOKE_QUERY = '''    
        MATCH (m:Mashup{{name:'{mashup_name}'}})
        MATCH (a) where id(a)={api_id}
        MERGE (m)-[:INVOKE]->(a)
        '''

GET_TYPE_BY_ID_QUERY = '''
        MATCH (n) WHERE id(n)={node_id}
        RETURN (labels(n))
'''

GET_IN_NODES_BY_ID_QUERY = '''
        MATCH (to) WHERE id(to)={node_id}
        MATCH (fr)-[]->(to) 
        RETURN id(fr)
'''

GET_API_IDS_QUERY = '''
        MATCH (api:API)
        RETURN id(api)
'''

GET_MASHUP_IDS_QUERY = '''
        MATCH (mashup:Mashup)
        RETURN id(mashup)
'''

fmt_q = '''match (n:Format)
return n
'''

class MashupKnowledgeGraph:

    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
        self.data_types = ['Mashup', 'API', 'Category', 'LatentTopic', 'User', 'Architecture', 'Format']
        for data_type in self.data_types:
            self.create_name_constraint(data_type)
    
    def close(self):
        self.driver.close()

    def create_name_constraint(self, data_type):
        with self.driver.session() as session:
            try:
                session.run(CREATE_CONSTRAINT_QUERY.format(data_type=data_type, unique_attribute='name')).data()
            except BaseException:
                return 

    def create_node(self, data_type, data_name):
        with self.driver.session() as session:
            try:
                session.run(CREATE_NODE_QUERY.format(data_type=data_type, data_name=data_name)).data()
            except BaseException:
                return 
    
    def get_node_id(self, data_type, data_name):
        with self.driver.session() as session:
            try:
                return list(session.run(GET_NODE_QUERY.format(data_type=data_type, data_name=data_name)).data()[0].values())[0]
            except BaseException:
                return None

    def remove_api_invoke(self, mashup_name): # 给定mashup名字，删除其api调用边，同时返回api的id list
        with self.driver.session() as session:
            try:
                msg = session.run(REMOVE_API_INVOKE_QUERY.format(mashup_name=mashup_name)).data()
                res = []
                for dic in msg:
                    res.extend(list(dic.values()))
                return res
            except BaseException:
                return None

    def recover_api_invoke(self, mashup_name, api_id_list): # 恢复mashup的api调用
        with self.driver.session() as session:
            for api_id in api_id_list:
                try:
                    session.run(RECOVER_API_INVOKE_QUERY.format(mashup_name=mashup_name, api_id=api_id))
                except BaseException:
                    return
    
    def create_relationship(self, from_type, from_name, to_type, to_name, relation_type):
        with self.driver.session() as session:
            try:
                session.run(CREATE_RELATIONSHIP_QUERY.format(from_type=from_type, from_name=from_name, to_type=to_type, to_name=to_name, relation_type=relation_type)).data() 
            except BaseException:
                return 

    def rand_walk(self, start_node_id, steps_count, walks_count, is_reverse):
        with self.driver.session() as session:
            paths = []
            try:
                for dic in session.run(RANDWALK_QUERY.format(
                                    start=start_node_id, 
                                    steps=steps_count, 
                                    walks=walks_count, 
                                    reverse='REVERSE' if is_reverse else 'NATURAL')
                                ).data():
                    print(dic)
                    paths.append(list(dic.values())[0])
            except BaseException:
                return [] 
            return paths
    
    def get_next(self, start_node_id)->list:
        with self.driver.session() as session:
            try:
                return [list(item.values())[0] for item in session.run(GET_NEXT_QUERY.format(id=start_node_id)).data()]
            except BaseException:
                return []

    def get_indegree(self, node_id)->int:
        with self.driver.session() as session:
            try:
                return list(session.run(GET_INDEGREE_QUERY.format(id=node_id)).data()[0].values())[0]
            except BaseException:
                return 0
    
    def get_type_by_id(self, node_id):
        with self.driver.session() as session:
            try:
                return list(session.run(GET_TYPE_BY_ID_QUERY.format(node_id=node_id)).data()[0].values())[0][0]
            except BaseException:
                return None

    def get_in_node_by_id(self, node_id):
        with self.driver.session() as session:
            try:
                return [list(dic.values())[0] for dic in session.run(GET_IN_NODES_BY_ID_QUERY.format(node_id=node_id)).data()]
            except BaseException:
                return None

    def get_api_ids(self):
        with self.driver.session() as session:
            try:
                return [list(dic.values())[0] for dic in session.run(GET_API_IDS_QUERY).data()]
            except BaseException:
                return None

    def get_mashup_ids(self):
        with self.driver.session() as session:
            try:
                return [list(dic.values())[0] for dic in session.run(GET_MASHUP_IDS_QUERY).data()]
            except BaseException:
                return None

    def fmt_test(self):
        with self.driver.session() as session:
            try:
                return session.run(fmt_q).data()
            except BaseException:
                return None

    def probe(self, path, i, sqrt_c)->dict:
        # print('probe',i,'path=',path)
        scores = {path[i]: 1.0}
        for j in range(i): # step
            temp_scores = {}
            for id, score in scores.items(): # current nodes
                nexts = self.get_next(id) # next nodes
                # print('cur',id,'nexts',nexts)
                for next in nexts:
                    if next != path[i - j - 1]:
                        if next in temp_scores:
                            temp_scores[next] = temp_scores[next] + sqrt_c * self.get_indegree(next) * score
                        else:
                            temp_scores[next] = sqrt_c / self.get_indegree(next) * score
            scores = temp_scores
        return scores

    def probe_sim(self, node_id, topk=10, c=0.6, r=1, l=1)->dict:
        sqrt_c = math.sqrt(c)
        topk_candidates = {}
        print(node_id)
        paths = self.rand_walk(start_node_id=node_id, steps_count=l, walks_count=r, is_reverse=True)
        print(paths)
        for s in range(r): # sample r times
            path = paths[s]
            # print(path)
            temp_candidates = {}
            for i in range(1, l + 1):
                scores = self.probe(path, i, sqrt_c)
                # print(scores)
                for id, score in scores.items():
                    if i == l and id not in topk_candidates:
                        topk_candidates[id] = 0.0
                    if id in temp_candidates:
                        temp_candidates[id] = temp_candidates[id] + score
                    else:
                        temp_candidates[id] = score
            # print('tc',temp_candidates)
            for v, score in topk_candidates.items():
                if v in temp_candidates:
                    topk_candidates[v] = score * s / (s + 1) + temp_candidates[v] / (s + 1)
                else:
                    topk_candidates[v] = score * s / (s + 1)
        res = {}
        for k, v in topk_candidates.items():
            if self.get_type_by_id(node_id=k) == 'Mashup':
                res[k] = v
        top_sims = sorted(res.items(), key=lambda item: -item[1])
        if len(top_sims) > topk:
            top_sims = top_sims[:topk]
        return top_sims

    def create_single_type_nodes(self, data_type, node_names):
        print('create nodes, type=', data_type, 'count=', len(node_names))
        for node_name in node_names:
            self.create_node(data_type=data_type, data_name=node_name)

    def create_single_type_relationships(self, dic, from_type, to_type, relation_type):
        for f, ts in dic.items():
            for t in ts:
                self.create_relationship(from_type=from_type, from_name=f, to_type=to_type, to_name=t, relation_type=relation_type)
            


        

if __name__ == "__main__":
    kg = MashupKnowledgeGraph("neo4j://localhost:7687", "neo4j", "hahaha")
    # print(kg.rand_walk(0, 3, 3, is_reverse=True))
    # print(kg.rand_walk(0, 3, 1, is_reverse=False))
    # print(kg.get_next(0))
    # print(kg.get_indegree(0))
    # print(kg.probe([0, 3], 1, math.sqrt(0.6)))
    # print(kg.probe_sim(node_id=0, topk=10, c=0.6, r=10, l=2))

    # kg.create_node(data_type='ft', data_name='fn')
    # kg.create_node(data_type='tt', data_name='tn')
    # kg.create_relationship(from_type='ft', from_name='fn', to_type='tt', to_name='tn', relation_type='RT')

    # kg.create_single_type_nodes('data_type', ['x','a','xxx'])

    # apis = ['a1','a2']
    # mashups = ['m1']
    # categories = ['c1', 'c2']
    # users = ['u1', 'u2']
    # latent_topics = ['lt1', 'lt2']
    # archs = ['arc1']

    # mashup_invoke_api = {'m1':['a1']}
    # api_usedby_mashup = {'a1':['m1']}
    # category_usedin_api = {'c1':['a1', 'a2'], 'c2':['a2']}
    # category_usedin_mashup = {'c1':['m1'], 'c2':['m1']}
    # user_follow_api = {'u1':['a1','a2']}
    # user_follow_mashup = {'u1':['m1']}
    # user_follow_user = {'u2':['u1']}
    # latenttopic_appearedin_mashup = {'lt1':['m1'], 'lt2':['m1']}
    # arch_support_api = {'arc1':['a1', 'a2']}

    # kg.create_single_type_nodes(data_type='API', node_names=apis)
    # kg.create_single_type_nodes(data_type='Mashup', node_names=mashups)
    # kg.create_single_type_nodes(data_type='Category', node_names=categories)
    # kg.create_single_type_nodes(data_type='User', node_names=users)
    # kg.create_single_type_nodes(data_type='LatentTopic', node_names=latent_topics)
    # kg.create_single_type_nodes(data_type='Architecture', node_names=archs)

    # kg.create_single_type_relationships(dic=mashup_invoke_api, from_type='Mashup', to_type='API', relation_type='INVOKE')
    # kg.create_single_type_relationships(dic=api_usedby_mashup, from_type='API', to_type='Mashup', relation_type='USEDBY')
    # kg.create_single_type_relationships(dic=category_usedin_api, from_type='Category', to_type='API', relation_type='USEDIN')
    # kg.create_single_type_relationships(dic=category_usedin_mashup, from_type='Category', to_type='Mashup', relation_type='USEDIN')
    # kg.create_single_type_relationships(dic=user_follow_api, from_type='User', to_type='API', relation_type='FOLLOW')
    # kg.create_single_type_relationships(dic=user_follow_mashup, from_type='User', to_type='Mashup', relation_type='FOLLOW')
    # kg.create_single_type_relationships(dic=user_follow_user, from_type='User', to_type='User', relation_type='FOLLOW')
    # kg.create_single_type_relationships(dic=latenttopic_appearedin_mashup, from_type='LatentTopic', to_type='Mashup', relation_type='APPEAREDIN')
    # kg.create_single_type_relationships(dic=arch_support_api, from_type='Architecture', to_type='API', relation_type='SUPPORT')


    # mashup = ['renttoownquest', 'proprover', 'betterhome', '222acres.com', 'propertywizza', 'globaround']
    mashup = ['betterhome', 'globaround']
    ma = mashups.Mashups()
    mashup = list(ma.mashup_dict.keys())
    wa = web_apis.Web_APIs()
    ma.train_topic_model()
    us = user.User()

    api = []
    mashup_to_api = {}
    api_to_mashup = {}
    for i in mashup:
        temp = []
        for j in ma.mashup_dict[i]['invoke_api']:
            if j in wa.api_dict:
                api.append(j)
                temp.append(j)
        mashup_to_api[i] = temp
    api = list(set(api))
    for i in mashup_to_api:
        for j in mashup_to_api[i]:
            temp = api_to_mashup[j] if j in api_to_mashup else []
            temp.append(i)
            api_to_mashup[j] = temp
    # print(api)
    # print(mashup_to_api)

    latent_topic = []
    latent_topic_to_mashup = {}
    for i in mashup:
        for j in ma.mashup_dict[i]['latent_topic']:
            latent_topic.append(j)
            temp = [] if j not in latent_topic_to_mashup else latent_topic_to_mashup[j]
            temp.append(i)
            latent_topic_to_mashup[j] = temp

    latent_topic = list(set(latent_topic))
    # print(latent_topic)

    category = []
    category_to_mashup = {}
    category_to_api = {}
    for i in mashup:
        for j in ma.mashup_dict[i]['category']:
            temp = [] if j not in category_to_mashup else category_to_mashup[j]
            temp.append(i)
            category_to_mashup[j] = temp
            category.append(j)
    for i in api:
        if i in wa.api_dict:
            for j in wa.api_dict[i]['category']:
                temp = [] if j not in category_to_api else category_to_api[j]
                temp.append(i)
                category_to_api[j] = temp
                category.append(j)
    category = list(set(category))
    # print(category)
    # print(category_to_mashup)
    # print(category_to_api)

    user = []
    user_to_mashup = {}
    user_to_api = {}
    for i in mashup:
        for j in ma.mashup_dict[i]['user']:
            if j in us.user_heat:
                temp = [] if j not in user_to_mashup else user_to_mashup[j]
                temp.append(i)
                user_to_mashup[j] = temp
                user.append(j)
    for i in api:
        if i in wa.api_dict:
            for j in wa.api_dict[i]['user']:
                if j in us.user_heat:
                    temp = [] if j not in user_to_api else user_to_api[j]
                    temp.append(i)
                    user_to_api[j] = temp
                    user.append(j)
    user = list(set(user))
    # print(user)
    user_to_user = {}
    for to in user:
        for fr in us.user_dict[to]:
            if fr in user:
                temp = [] if fr not in user_to_user else user_to_user[fr]
                temp.append(to)
                user_to_user[fr] = temp
    # print(user_to_user)

    arch = []
    arch_to_api = {}
    for i in api:
        if i in wa.api_dict and 'architecture' in wa.api_dict[i]:
            for j in wa.api_dict[i]['architecture']:
                arch.append(j)
                temp = [] if j not in arch_to_api else arch_to_api[j]
                temp.append(i)
                arch_to_api[j] = temp
    arch = list(set(arch))
    # print(arch)
    # print(arch_to_api)

    fmt = []
    request_format_to_api = {}
    for i in api:
        if i in wa.api_dict and 'request_format' in wa.api_dict[i]:
            for j in wa.api_dict[i]['request_format']:
                fmt.append(j)
                temp = [] if j not in request_format_to_api else request_format_to_api[j]
                temp.append(i)
                request_format_to_api[j] = temp
    # print(request_format)

    response_format_to_api = {}
    for i in api:
        if i in wa.api_dict and 'response_format' in wa.api_dict[i]:
            for j in wa.api_dict[i]['response_format']:
                fmt.append(j)
                temp = [] if j not in response_format_to_api else response_format_to_api[j]
                temp.append(i)
                response_format_to_api[j] = temp
    fmt = list(set(fmt))
    # print(response_format)

    kg.create_single_type_nodes(data_type='Mashup', node_names=mashup)
    kg.create_single_type_nodes(data_type='API', node_names=api)
    kg.create_single_type_nodes(data_type='LatentTopic', node_names=latent_topic)
    kg.create_single_type_nodes(data_type='Category', node_names=category)
    kg.create_single_type_nodes(data_type='User', node_names=user)
    kg.create_single_type_nodes(data_type='Architecture', node_names=arch)
    kg.create_single_type_nodes(data_type='Format', node_names=fmt)

    kg.create_single_type_relationships(dic=mashup_to_api, from_type='Mashup', to_type='API', relation_type='INVOKE')
    kg.create_single_type_relationships(dic=api_to_mashup, from_type='API', to_type='Mashup', relation_type='COMPOSED')

    kg.create_single_type_relationships(dic=latent_topic_to_mashup, from_type='LatentTopic', to_type='Mashup', relation_type='PARTOF')

    kg.create_single_type_relationships(dic=category_to_mashup, from_type='Category', to_type='Mashup', relation_type='PARTOF')
    kg.create_single_type_relationships(dic=category_to_api, from_type='Category', to_type='API', relation_type='PARTOF')

    kg.create_single_type_relationships(dic=user_to_mashup, from_type='User', to_type='Mashup', relation_type='FOLLOW')
    kg.create_single_type_relationships(dic=user_to_api, from_type='User', to_type='API', relation_type='FOLLOW')
    kg.create_single_type_relationships(dic=user_to_user, from_type='User', to_type='User', relation_type='FOLLOW')

    kg.create_single_type_relationships(dic=arch_to_api, from_type='Architecture', to_type='API', relation_type='SUPPORTED_ARCH')
    kg.create_single_type_relationships(dic=request_format_to_api, from_type='Format', to_type='API', relation_type='SUPPORTED_REQ')
    kg.create_single_type_relationships(dic=response_format_to_api, from_type='Format', to_type='API', relation_type='SUPPORTED_RESP')

    kg.close()
